'use client';

import type { LanguageModelUsage } from 'ai';
import { type ComponentProps, createContext, useContext } from 'react';
import { estimateCost } from 'tokenlens';
import { Button } from '../../components/button';
import { HoverCard, HoverCardContent, HoverCardTrigger } from '../../components/hover-card';
import { Progress } from '../../components/progress';
import { cn } from '../../utils';

const PERCENT_MAX = 100;
const clamp01 = (v: number) => Math.min(1, Math.max(0, v));
const pct = (used: number, max: number) => (max > 0 ? clamp01(used / max) : 0);
const ICON_RADIUS = 10;
const ICON_VIEWBOX = 24;
const ICON_CENTER = 12;
const ICON_STROKE_WIDTH = 2;

type ContextSchema = {
    usedTokens: number;
    maxTokens: number;
    usage?: LanguageModelUsage;
    modelId?: string;
};

const ContextContext = createContext<ContextSchema | null>(null);

const useContextValue = () => {
    const context = useContext(ContextContext);

    if (!context) {
        throw new Error('Context components must be used within Context');
    }

    return context;
};

export type ContextProps = ComponentProps<typeof HoverCard> & ContextSchema;

export const Context = ({ usedTokens, maxTokens, usage, modelId, ...props }: ContextProps) => (
    <ContextContext.Provider
        value={{
            usedTokens,
            maxTokens,
            usage,
            modelId,
        }}
    >
        <HoverCard closeDelay={0} openDelay={0} {...props} />
    </ContextContext.Provider>
);

const ContextIcon = () => {
    const { usedTokens, maxTokens } = useContextValue();
    const circumference = 2 * Math.PI * ICON_RADIUS;
    const usedPercent = pct(usedTokens, maxTokens);
    const dashOffset = circumference * (1 - usedPercent);

    return (
        <svg
            aria-label="Model context usage"
            height="20"
            role="img"
            style={{ color: 'currentcolor' }}
            viewBox={`0 0 ${ICON_VIEWBOX} ${ICON_VIEWBOX}`}
            width="20"
        >
            <circle
                cx={ICON_CENTER}
                cy={ICON_CENTER}
                fill="none"
                opacity="0.25"
                r={ICON_RADIUS}
                stroke="currentColor"
                strokeWidth={ICON_STROKE_WIDTH}
            />
            <circle
                cx={ICON_CENTER}
                cy={ICON_CENTER}
                fill="none"
                opacity="0.7"
                r={ICON_RADIUS}
                stroke="currentColor"
                strokeDasharray={`${circumference} ${circumference}`}
                strokeDashoffset={dashOffset}
                strokeLinecap="round"
                strokeWidth={ICON_STROKE_WIDTH}
                style={{ transformOrigin: 'center', transform: 'rotate(-90deg)' }}
            />
        </svg>
    );
};

export type ContextTriggerProps = ComponentProps<typeof Button>;

export const ContextTrigger = ({ children, ...props }: ContextTriggerProps) => {
    const { usedTokens, maxTokens } = useContextValue();
    const usedPercent = pct(usedTokens, maxTokens);
    const renderedPercent = new Intl.NumberFormat('en-US', {
        style: 'percent',
        maximumFractionDigits: 0,
    }).format(usedPercent);

    return (
        <HoverCardTrigger asChild>
            {children ?? (
                <Button type="button" variant="ghost" {...props}>
                    <ContextIcon />
                    <span className="text-xs font-medium text-muted-foreground">
                        {renderedPercent}
                    </span>
                </Button>
            )}
        </HoverCardTrigger>
    );
};

export type ContextContentProps = ComponentProps<typeof HoverCardContent>;

export const ContextContent = ({ className, ...props }: ContextContentProps) => (
    <HoverCardContent
        className={cn('min-w-[240px] divide-y overflow-hidden p-0', className)}
        {...props}
    />
);

export type ContextContentHeader = ComponentProps<'div'>;

export const ContextContentHeader = ({ children, className, ...props }: ContextContentHeader) => {
    const { usedTokens, maxTokens } = useContextValue();
    const usedPercent = pct(usedTokens, maxTokens);
    const displayPct = new Intl.NumberFormat('en-US', {
        style: 'percent',
        maximumFractionDigits: 1,
    }).format(usedPercent);
    const used = new Intl.NumberFormat('en-US', {
        notation: 'compact',
    }).format(usedTokens);
    const total = new Intl.NumberFormat('en-US', {
        notation: 'compact',
    }).format(maxTokens);

    return (
        <div className={cn('w-full space-y-2 p-3', className)} {...props}>
            {children ?? (
                <>
                    <div className="flex items-center justify-between gap-3 text-xs">
                        <p>{displayPct}</p>
                        <p className="font-mono text-muted-foreground">
                            {used} / {total}
                        </p>
                    </div>
                    <div className="space-y-2">
                        <Progress
                            className="bg-muted"
                            value={Math.round(usedPercent * PERCENT_MAX)}
                        />
                    </div>
                </>
            )}
        </div>
    );
};

export type ContextContentBody = ComponentProps<'div'>;

export const ContextContentBody = ({ children, className, ...props }: ContextContentBody) => (
    <div className={cn('w-full p-3', className)} {...props}>
        {children}
    </div>
);

export type ContextContentFooter = ComponentProps<'div'>;

export const ContextContentFooter = ({ children, className, ...props }: ContextContentFooter) => {
    const { modelId, usage } = useContextValue();
    const costUSD = modelId
        ? estimateCost({
              modelId,
              usage: {
                  input: usage?.inputTokens ?? 0,
                  output: usage?.outputTokens ?? 0,
              },
          }).totalUSD
        : undefined;
    const totalCost = new Intl.NumberFormat('en-US', {
        style: 'currency',
        currency: 'USD',
    }).format(costUSD ?? 0);

    return (
        <div
            className={cn(
                'flex w-full items-center justify-between gap-3 bg-secondary p-3 text-xs',
                className,
            )}
            {...props}
        >
            {children ?? (
                <>
                    <span className="text-muted-foreground">Total cost</span>
                    <span>{totalCost}</span>
                </>
            )}
        </div>
    );
};

export type ContextInputUsageProps = ComponentProps<'div'>;

export const ContextInputUsage = ({ className, children, ...props }: ContextInputUsageProps) => {
    const { usage, modelId } = useContextValue();
    const inputTokens = usage?.inputTokens ?? 0;

    if (children) {
        return children;
    }

    if (!inputTokens) {
        return null;
    }

    const inputCost = modelId
        ? estimateCost({
              modelId,
              usage: { input: inputTokens, output: 0 },
          }).totalUSD
        : undefined;
    const inputCostText = new Intl.NumberFormat('en-US', {
        style: 'currency',
        currency: 'USD',
    }).format(inputCost ?? 0);

    return (
        <div className={cn('flex items-center justify-between text-xs', className)} {...props}>
            <span className="text-muted-foreground">Input</span>
            <TokensWithCost costText={inputCostText} tokens={inputTokens} />
        </div>
    );
};

export type ContextOutputUsageProps = ComponentProps<'div'>;

export const ContextOutputUsage = ({ className, children, ...props }: ContextOutputUsageProps) => {
    const { usage, modelId } = useContextValue();
    const outputTokens = usage?.outputTokens ?? 0;

    if (children) {
        return children;
    }

    if (!outputTokens) {
        return null;
    }

    const outputCost = modelId
        ? estimateCost({
              modelId,
              usage: { input: 0, output: outputTokens },
          }).totalUSD
        : undefined;
    const outputCostText = new Intl.NumberFormat('en-US', {
        style: 'currency',
        currency: 'USD',
    }).format(outputCost ?? 0);

    return (
        <div className={cn('flex items-center justify-between text-xs', className)} {...props}>
            <span className="text-muted-foreground">Output</span>
            <TokensWithCost costText={outputCostText} tokens={outputTokens} />
        </div>
    );
};

export type ContextReasoningUsageProps = ComponentProps<'div'>;

export const ContextReasoningUsage = ({
    className,
    children,
    ...props
}: ContextReasoningUsageProps) => {
    const { usage, modelId } = useContextValue();
    const reasoningTokens = usage?.reasoningTokens ?? 0;

    if (children) {
        return children;
    }

    if (!reasoningTokens) {
        return null;
    }

    const reasoningCost = modelId
        ? estimateCost({
              modelId,
              usage: { reasoningTokens },
          }).totalUSD
        : undefined;
    const reasoningCostText = new Intl.NumberFormat('en-US', {
        style: 'currency',
        currency: 'USD',
    }).format(reasoningCost ?? 0);

    return (
        <div className={cn('flex items-center justify-between text-xs', className)} {...props}>
            <span className="text-muted-foreground">Reasoning</span>
            <TokensWithCost costText={reasoningCostText} tokens={reasoningTokens} />
        </div>
    );
};

export type ContextCacheUsageProps = ComponentProps<'div'>;

export const ContextCacheUsage = ({ className, children, ...props }: ContextCacheUsageProps) => {
    const { usage, modelId } = useContextValue();
    const cacheTokens = usage?.cachedInputTokens ?? 0;

    if (children) {
        return children;
    }

    if (!cacheTokens) {
        return null;
    }

    const cacheCost = modelId
        ? estimateCost({
              modelId,
              usage: { cacheReads: cacheTokens, input: 0, output: 0 },
          }).totalUSD
        : undefined;
    const cacheCostText = new Intl.NumberFormat('en-US', {
        style: 'currency',
        currency: 'USD',
    }).format(cacheCost ?? 0);

    return (
        <div className={cn('flex items-center justify-between text-xs', className)} {...props}>
            <span className="text-muted-foreground">Cache</span>
            <TokensWithCost costText={cacheCostText} tokens={cacheTokens} />
        </div>
    );
};

const TokensWithCost = ({
    tokens,
    costText,
    showCost = false,
}: {
    tokens?: number;
    costText?: string;
    showCost?: boolean;
}) => {
    return (
        <span>
            {tokens === undefined
                ? '—'
                : new Intl.NumberFormat('en-US', {
                      notation: 'compact',
                  }).format(tokens)}
            {showCost && costText ? (
                <span className="ml-2 text-muted-foreground">• {costText}</span>
            ) : null}
        </span>
    );
};
